import os
import numpy as np
import pickle
from PIL import Image
from matplotlib import pyplot as plt
import chineseize_matplotlib


def sigmoid(x):
    return 1 / (1 + np.exp(-x))


def sigmoid_derivative(x):
    return x * (1 - x)


def load_model(file_path):
    with open(file_path, 'rb') as file:
        model = pickle.load(file)
    return model


def predict(image, model):
    weights1 = model['weights1']
    weights2 = model['weights2']

    # 预处理输入图像
    image = image.reshape(1, -1) / 255.0  # 将图像展平并归一化

    # 前向传播
    hidden_layer_input = np.dot(image, weights1)
    hidden_layer_output = sigmoid(hidden_layer_input)
    output_layer_input = np.dot(hidden_layer_output, weights2)
    predicted_output = sigmoid(output_layer_input)

    # 返回最可能的类别,onehot编码中最大的
    return np.argmax(predicted_output, axis=1)


def load_and_resize_images(folder_path, target_size=(128, 128)):
    resized_images = []
    for file in os.listdir(folder_path):
        if file.endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
            img_path = os.path.join(folder_path, file)
            with Image.open(img_path) as img:
                img = img.convert('RGB')
                img_resized = img.resize(target_size)
                resized_images.append(np.array(img_resized))

    return resized_images


# 以上部分和训练.py一样

# 加载模型,我把模型保存为二进制文件,就是为了在不同的地方都能用模型
model = load_model('model.pkl')

path = './traffic_light_data/val'
predict_image = load_and_resize_images(path, target_size=(32, 32))

path_name = []


def one_hot(array):
    for i in range(len(array)):
        if array[i] == 1:
            return i


red_num = 0
green_num = 0
yellow_num = 0

accuracies = []
correct_predictions = 0
# 模型文件与结果一样的
for file in os.listdir(path):
    if file.startswith('green'):
        green_num += 1
    elif file.startswith('red'):
        red_num += 1
    elif file.startswith('yellow'):
        yellow_num += 1

true_labels = np.concatenate([np.tile([0, 1, 0], (green_num, 1)),  # 绿灯
                              np.tile([1, 0, 0], (red_num, 1)),  # 红灯
                              np.tile([0, 0, 1], (yellow_num, 1))])  # 黄灯

class_names = ['红', '绿', '黄']

# 直接预测,没有封装函数

for idx in range(len(predict_image)):
    prediction = predict(predict_image[idx], model)
    if one_hot(true_labels[idx]) != prediction[0]:
        print("预测的不对")
        print("\n推测的结果为:", class_names[prediction[0]])
        print("实际为:", class_names[one_hot(true_labels[idx])], '\n')
    if one_hot(true_labels[idx]) == prediction[0]:
        correct_predictions += 1
    accuracy = correct_predictions / (idx + 1)
    accuracies.append(accuracy)
    print("\n推测的结果为:", class_names[prediction[0]])
    print("实际为:", class_names[one_hot(true_labels[idx])], '\n')

# 绘制准确率图线
plt.plot(accuracies)
plt.title('随图像数量变化的准确率')
plt.xlabel('图像数量')
plt.ylabel('准确率')
plt.grid(True)
plt.savefig('准确率.png')  # 保存图像
plt.show()  # 显示图像
